Skip to content

Enable range learning for QAT #2033

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 29, 2025
Merged

Enable range learning for QAT #2033

merged 1 commit into from
May 29, 2025

Conversation

andrewor14
Copy link
Contributor

@andrewor14 andrewor14 commented Apr 9, 2025

Summary: This commit adds the option for QAT users to use range learning during training. Range learning means we train the scale and zero point instead of recomputing them based on the input in every iteration.

Example usage:

import torch
from torchao.quantization import quantize_
from torchao.quantization.qat import (
    FakeQuantizeConfig,
    IntXQuantizationAwareTrainingConfig,
    initialize_fake_quantizers,
)

config = FakeQuantizeConfig(
    torch.int8,
    "per_channel",
    is_dynamic=False,
    range_learning=True,
    scale_precision=torch.float32,
    zero_point_precision=torch.float32,
)
m = M()
example_inputs = (torch.randn(16, 32),)
quantize_(m, IntXQuantizationAwareTrainingConfig(weight_config=config))

# New required step to turn scales and zero points into trainable
# `nn.Parameters`, must be called before initializing the optimizer
initialize_fake_quantizers(m, example_inputs)

# initialize the optimizer
# do training

Test Plan:
python test/quantization/test_qat.py -k test_fake_quantize_config_dynamic_and_range_learning
python test/quantization/test_qat.py -k test_fake_quantizer_range_learning
python test/quantization/test_qat.py -k test_qat_range_learning

Copy link

pytorch-bot bot commented Apr 9, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2033

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 045b99e with merge base 0aa8dbd (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 9, 2025
@andrewor14 andrewor14 added topic: new feature Use this tag if this PR adds a new feature and removed CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. labels Apr 9, 2025
@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 9, 2025
@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@andrewor14 andrewor14 force-pushed the range-learning branch 2 times, most recently from 912a925 to cb98579 Compare May 1, 2025 22:02
@drisspg
Copy link
Contributor

drisspg commented May 2, 2025

I think this needs a rebase

@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

**Summary:** This commit adds the option for QAT users to use range
learning during training. Range learning means we train the scale
and zero point instead of recomputing them based on the input
at every iteration.

Example usage:
```
import torch
from torchao.quantization import quantize_
from torchao.quantization.qat import (
    FakeQuantizeConfig,
    IntXQuantizationAwareTrainingConfig,
    initialize_fake_quantizers,
)

config = FakeQuantizeConfig(
    torch.int8,
    "per_channel",
    is_dynamic=False,
    range_learning=True,
    scale_precision=torch.float32,
    zero_point_precision=torch.float32,
)
m = M()
example_inputs = (torch.randn(16, 32),)
quantize_(m, IntXQuantizationAwareTrainingConfig(weight_config=config))

# New required step to turn scales and zero points into trainable
# `nn.Parameters`, must be called before initializing the optimizer
initialize_fake_quantizers(m, example_inputs)

# initialize the optimizer
# do training
```

**Test Plan:**
python test/quantization/test_qat.py -k test_fake_quantize_config_dynamic_and_range_learning
python test/quantization/test_qat.py -k test_fake_quantizer_range_learning
python test/quantization/test_qat.py -k test_qat_range_learning
@facebook-github-bot
Copy link
Contributor

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@facebook-github-bot facebook-github-bot merged commit 1b61c82 into main May 29, 2025
20 of 21 checks passed
@andrewor14 andrewor14 mentioned this pull request May 29, 2025
4 tasks
@gau-nernst
Copy link
Collaborator

@andrewor14 Did you downgrade cutlass to 3.8 by mistake in this PR? I'm trying to build from source and it fails to build the MX kernels.

@andrewor14
Copy link
Contributor Author

Ah my mistake, let me look into a fix

andrewor14 added a commit that referenced this pull request May 30, 2025
Cutlass was accidentally downgraded to 3.8 in
#2033. This commit undoes
this change and brings cutlass back to the original 3.9 commit.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: new feature Use this tag if this PR adds a new feature
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants